import argparse
import os
import re
import time
from time import sleep
from datasets import audio
import tensorflow as tf
from hparams import hparams, hparams_debug_string
from infolog import log
from tacotron.synthesizer import Synthesizer
from tqdm import tqdm


def generate_fast(model, text):
	# out_dir = "/opt/tiger/arnold_test/workstation/data/single_speaker/logs-yinsu/synthesis/"
	# log_dir = "/opt/tiger/arnold_test/workstation/data/single_speaker/logs-yinsu/synthesis/"

	out_dir = "/opt/tiger/arnold_test/workstation/data/single_speaker/logs-yinsu/synthesis/"
	log_dir = "/opt/tiger/arnold_test/workstation/data/single_speaker/logs-yinsu/synthesis/"
	model.synthesize(text, None, out_dir, log_dir)


def run_live(args, checkpoint_path, hparams):
	#Log to Terminal without keeping any records in files
	log(hparams_debug_string())
	synth = Synthesizer()
	synth.load(checkpoint_path, hparams)

	#Generate fast greeting message
	# greetings = 'Hello, Welcome to the Live testing tool. Please type a message and I will try to read it!'
	# log(greetings)
	# generate_fast(synth, greetings)
    
	
	# text = ['sil a4 n ii3 k e3 zh en1 uei3 d a4 ia4 sil EOS']
	# text = ['k uai4 d ian3 g an4 b a4 sil sil']
	# text = ['sil zh e4 x ia4 uan2 l e5 sil EOS']
	# text = ['sil a4 x ia4 v3 l e5 sil EOS']
	# text = ['sil t ai4 b ang4 l e5 sil EOS']
	# text = ['sil uo3 zh en1 d e5 i3 uei2 n i3 sh i4 zh e4 g e4 i4 s ii5 sil']
	# text = ['sil uo3 z ai4 l uen4 uen2 sh ang4 k an4 d ao4 n i3 d e5 m ing2 z ii4 l e5 sil']
	# text = ['sil ei1 s ei1 m i3 l an2 ing2 q iou2 l e5 sil']
	# text = ['sil uo2 zh e4 c ii4 k ao3 sh ii4 g ang1 g ang1 t ong1 g uo4 sil']
	# text = ['sil j in1 t ian1 sh ii4 x ing1 q i1 t ian1 sil EOS']
	# text = ['sil sil ai4 n i3 zh e4 r en2 sil']
	# text = ['sil d ian4 h ua4 l ing2 x iang3 l e5 sil']
	# text = ['sil t a1 j iou4 k uai4 l ai2 l e5 sil']
	# text = ['sil l u4 sh ang4 r en2 zh en1 d uo1 a4 sil']
	text = ['sil m ing2 t ian1 uo3 iao4 b an1 j ia1 l e5 sil']
	# text = ['sil zh e4 j ian4 sh ii4 sh ii4 t a1 g an4 d e5 sil']
	# text = ['sil n i3 zh e4 d uan4 sh ii2 j ian1 b ian4 sh ou4 l e5 sil']
	# text = ['sil g uo4 l iang3 t ian1 x ve2 x iao4 j iou4 iao4 k ai1 x ve2 l e5 sil']
	# text = ['sil z uo2 t ian1 uan3 sh ang4 uo3 z uo4 l e5 i2 g e4 m eng4 sil EOS']
	# text = ['sil iou3 i2 l iang4 ch e1 x iang4 uo3 m en2 k ai1 l e5 g uo4 l ai2 sil EOS']

	generate_fast(synth, text)
	#Interaction loop
	# while True:
	# 	try:
	# 		# text = input()
	# 		text = 'ta1 shi4 yin1 pin2 ling3 yu4 de tao2 bao3 tian1 mao1'
	# 		generate_fast(synth, text)

	# 	except KeyboardInterrupt:
	# 		leave = 'Thank you for testing our features. see you soon.'
	# 		log(leave)
	# 		generate_fast(synth, leave)
	# 		sleep(2)
	# 		break

def run_eval(args, checkpoint_path, output_dir, hparams, sentences):
	eval_dir = os.path.join(output_dir, 'eval')
	log_dir = os.path.join(output_dir, 'logs-eval')

	#Create output path if it doesn't exist
	os.makedirs(eval_dir, exist_ok=True)
	os.makedirs(log_dir, exist_ok=True)
	os.makedirs(os.path.join(log_dir, 'wavs'), exist_ok=True)
	os.makedirs(os.path.join(log_dir, 'plots'), exist_ok=True)

	log(hparams_debug_string())
	synth = Synthesizer()
	synth.load(checkpoint_path, hparams)

	delta_size = hparams.tacotron_synthesis_batch_size if hparams.tacotron_synthesis_batch_size < len(sentences) else len(sentences)
	batch_sentences = [sentences[i: i+hparams.tacotron_synthesis_batch_size] for i in range(0, len(sentences), delta_size)]

	start = time.time()
	for i, batch in enumerate(tqdm(batch_sentences)):
		audio.save_wav(synth.eval(batch), os.path.join(log_dir, 'wavs', 'eval_batch_{:03}.wav'.format(i)), hparams)
	log('\nGenerated batches of size {} in {:.3f} sec'.format(delta_size, time.time() - start))

	return eval_dir


def tacotron_synthesize(args, hparams, checkpoint, sentences=None):
	output_dir = 'tacotron_' + args.output_dir

	try:
		checkpoint_path = tf.train.get_checkpoint_state(checkpoint).model_checkpoint_path
		log('loaded model at {}'.format(checkpoint_path))
	except:
		raise RuntimeError('Failed to load checkpoint at {}'.format(checkpoint))

	if args.mode == 'eval':
		return run_eval(args, checkpoint_path, output_dir, hparams, sentences)
	else:
		run_live(args, checkpoint_path, hparams)
